#include "common.h"
#define MAX_EVENT 5

int server = 0;
int port = 8000;
char hostname[32];

void server_routine();
void client_routine();
int onRead(int, int, int);
void run(int);

int main(int argc, char* argv[]) {
	Args args = parsecmdline(argc, argv);
	if (args.empty() || args.find("h") == args.end()
			|| (args.find("l") != args.end() && args.find("p") != args.end())
			|| (args.find("l") == args.end() && args.find("p") == args.end())) {
		printf("Usage:\n  %s <-h hostname> <-p port or -l port>\n", argv[0]);
		return 1;
	}

  
	strcpy(hostname, args["h"].c_str());

	if (args.find("l") != args.end()) {
		port = atoi(args["l"].c_str());
		server = 1;
	}
	else {
		port = atoi(args["p"].c_str());
	}


	if (server) {
		server_routine();
	}
	else {
		client_routine();
	}
	return 0;
}

void server_routine() {
  int listenfd, clientfd, ret, len, epfd;
	int *fd;
	pthread_t th;
  struct sockaddr_in servaddr, cliaddr;
	socklen_t servaddrlen, cliaddrlen;

	servaddr = resolve(hostname, port);

	listenfd = socket(AF_INET, SOCK_STREAM, 0);
	if (listenfd < 0) ERR_EXIT("socket");

	ret  = bind(listenfd, (struct sockaddr*)&servaddr, sizeof(servaddr));
	if (ret < 0) ERR_EXIT("bind");

	ret = listen(listenfd, 5);
	if (ret < 0) ERR_EXIT("listen");


	cliaddrlen = sizeof(cliaddr);
	clientfd = accept(listenfd, (struct sockaddr*)&cliaddr, &cliaddrlen); 
	if (clientfd < 0) ERR_EXIT("accept");

	run(clientfd);
}

void client_routine() {
	int sockfd, ret, *fd;
  struct sockaddr_in servaddr;
	char *buf;
	pthread_t th;

	servaddr = resolve(hostname, port);
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd < 0) ERR_EXIT("socket");

	ret = connect(sockfd, (struct sockaddr*)&servaddr, sizeof(servaddr));
	if (ret < 0) ERR_EXIT("connect");
	
	run(sockfd);
}

int onRead(int epfd, int fd, int clientfd) {
	int nr, ret;
	char *buf;

	buf = (char*)malloc(8192);
	assert(buf);

	if (fd != STDIN_FILENO) {
		nr = read(fd, buf, 8192);
		if (nr < 0) ERR_EXIT("read");
		else if (nr == 0) {
			// I also have no data to write.
			puts("peer closed");
			shutdown(fd, SHUT_WR);
			ret = epoll_ctl(epfd, EPOLL_CTL_DEL, fd, NULL);
			if (ret < 0) ERR_EXIT("epoll_ctl in onRead");
			return -1;
		}
		else {
			write(STDOUT_FILENO, buf, nr);
			return 0;
		}
	}
	else {
		nr = read(fd, buf, 8192);
		if (nr < 0) ERR_EXIT("read");
		else if (nr == 0) {
			// I have no data to write.
			shutdown(clientfd, SHUT_WR);
			ret = epoll_ctl(epfd, EPOLL_CTL_DEL, STDIN_FILENO, NULL);
			if (ret < 0) ERR_EXIT("epoll_ctl in onRead");
			return -1;
		}
		else {
			write(clientfd, buf, nr);
			return 0;
		}
	}
}

void run(int sockfd) {
	int epfd, ret, i;
	struct epoll_event ev, evts[MAX_EVENT];
	epfd = epoll_create(MAX_EVENT);
	if (epfd < 0) ERR_EXIT("epoll_create");

	// add sockfd 
	ev.events = EPOLLIN;
	ev.data.fd = sockfd;

	ret = epoll_ctl(epfd, EPOLL_CTL_ADD, sockfd, &ev);
	if (ret < 0) ERR_EXIT("epoll_ctl add clientfd");

	// add STDIN_FILENO
	ev.events = EPOLLIN;
	ev.data.fd = STDIN_FILENO;

	ret = epoll_ctl(epfd, EPOLL_CTL_ADD, STDIN_FILENO, &ev);
	if (ret < 0) ERR_EXIT("epoll_ctl add STDIN_FILENO");

	while(1) {
		ret = epoll_wait(epfd, evts, MAX_EVENT, 0);
	  if (ret < 0) ERR_EXIT("epoll_wait");	
		for (i = 0; i < ret; ++i) {
			if (evts[i].events & EPOLLIN) {
				ret = onRead(epfd, evts[i].data.fd, sockfd);
				if (ret < 0) return;
			}
			if (evts[i].events & EPOLLHUP) {
				puts("EPOLL HUP");
				exit(1);
			}
			if (evts[i].events & EPOLLERR) {
				puts("EPOLL ERR");
				exit(1);
			}
		}
	}
}
